import torch
import torch.nn as nn
from transformers import BertModel

from bt_config import BT_Config

class BT_model(nn.Module):
    def __init__(self, config: BT_Config):
        super(BT_model, self).__init__()
        self.bert = BertModel.from_pretrained(config.bert_path)
        
        # 冻结 BERT 参数
        for param in self.bert.parameters():
            param.requires_grad = False
        
        self.fc = nn.Linear(config.hidden_size, 10)  # 10分类任务

    def forward(self, input_ids, attention_mask):
        with torch.no_grad():
            outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        x1 = outputs.pooler_output  # [CLS]位置的向量
        logits = self.fc(x1)
        return logits
